import numpy as np
import cv2
import os
from torch.utils.data import Dataset, DataLoader
import torch


class loader(Dataset):
    def __init__(self, path, root, header=True):
        self.lines = []
        if isinstance(path, list):
            for i in path:
                with open(i) as f:
                    line = f.readlines()
                    if header: line.pop(0)
                    self.lines.extend(line)
        else:
            with open(path) as f:
                self.lines = f.readlines()
                if header: self.lines.pop(0)

        self.root = root

    def __len__(self):
        return len(self.lines)

    def __getitem__(self, idx):

        line = self.lines[idx]
        line = line.strip().split(" ")

        gaze2d = line[1]
        label_t = np.array(gaze2d.split(",")).astype("float")
        temp = label_t[0]
        label_t[0] = label_t[1]
        label_t[1] = temp
        gaze_rem = label_t

        face = line[0]
        face = face.replace("\\", "/")
        fimg = cv2.imread(os.path.join(self.root, face)) / 255.0

        fimg = fimg.transpose(2, 0, 1)


        img = {
            "face": torch.from_numpy(fimg).type(torch.FloatTensor),
        }
        target = {
            "gaze_yaw_pitch": torch.from_numpy(gaze_rem).type(torch.FloatTensor),
            "img_name": os.path.join(self.root, face),
        }

        return img, target


def txtload(labelpath, imagepath, batch_size, shuffle=True, num_workers=0, header=True):
    # labelpath = ["/home/xian/Downloads/Gaze-Net-main/Dataset/FaceBased/ETH-Gaze/Label/train_del_face_3000.label"]
    # imagepath = "/home/xian/Downloads/Gaze-Net-main/Dataset/FaceBased/ETH-Gaze/Image"
    dataset = loader(labelpath, imagepath, header)
    print(f"[Read Data]: Total num: {len(dataset)}")
    print(f"[Read Data]: Label path: {labelpath}")
    load = DataLoader(dataset, batch_size=batch_size, shuffle=shuffle, num_workers=num_workers)
    return load

if __name__ == "__main__":
    path = '/home/ys/Downloads/Gaze-Net-main/Dataset/FaceBased/ETH-Gaze/Label/train_del.label'
    d = loader(path, "/home/ys/Downloads/Gaze-Net-main/Dataset/FaceBased/ETH-Gaze/Image")
    s_path = "/home/ys/Downloads/Gaze-Net-main/del_eth"
    print(len(d))

    for i in range(len(d)):
        (data, label) = d.__getitem__(i)
        target_reg = label['target_reg']
        yaw = target_reg[1]
        # pitch = target_reg[0]

        gaze_rem = yaw + 90
        gaze_rem = gaze_rem.round()
        if gaze_rem >= 180:
            gaze_rem = 179
        if gaze_rem <= 0:
            gaze_rem = 0
        #print(int(np.asarray(gaze_rem)))
        temp_class = int(np.asarray(gaze_rem))
        path = str(temp_class)
        file_old = label['img_name']
        face_img = cv2.imread(file_old)


